import os
import time
from dotenv import load_dotenv
from qcloud_hunyuan3d import QCloudHunyuan3D

# 加载环境变量
load_dotenv()

def text_to_3d_example(action_type=None):
    """示例：通过文本生成3D模型"""
    print("\n=== 示例：通过文本生成3D模型 ===")
    
    # 初始化客户端
    client = QCloudHunyuan3D()
    
    # 设置文本描述和可选参数
    prompt = "一匹马，高细节，真实风格"
    options = {}
    
    # 如果指定了action类型，则添加到options中
    if action_type:
        options['action'] = action_type
        print(f"使用接口类型: {action_type}")
    else:
        print("使用默认接口类型")
    
    try:
        print(f"正在生成3D模型，文本描述：{prompt}")
        # 调用API生成3D模型
        result = client.generate_3d_from_text(prompt, options)
        
        print("生成请求已提交，响应结果：")
        print(result)
        
        # 如果返回了任务ID，可以查询任务状态
        if 'Response' in result and 'JobId' in result['Response']:
            job_id = result['Response']['JobId']
            print(f"任务ID: {job_id}")
            
            # 查询任务状态
            check_task_status_example(client, job_id, action_type)
            
            # 监控任务直到完成
            # monitor_task_until_completion(client, job_id, action_type)
            
        return result
        
    except Exception as e:
        print(f"生成失败: {str(e)}")

def image_url_to_3d_example(action_type=None):
    """示例：通过图片URL生成3D模型"""
    print("\n=== 示例：通过图片URL生成3D模型 ===")
    
    # 初始化客户端
    client = QCloudHunyuan3D()
    
    # 图片URL（替换为实际的图片URL）
    image_url = "https://example.com/sample.jpg"
    options = {}
    
    # 如果指定了action类型，则添加到options中
    if action_type:
        options['action'] = action_type
        print(f"使用接口类型: {action_type}")
    else:
        print("使用默认接口类型")
    
    try:
        print(f"正在根据图片生成3D模型，图片URL：{image_url}")
        # 调用API生成3D模型
        result = client.generate_3d_from_image(image_url, options)
        
        print("生成请求已提交，响应结果：")
        print(result)
        
        # 如果返回了任务ID，可以查询任务状态
        if 'Response' in result and 'JobId' in result['Response']:
            job_id = result['Response']['JobId']
            print(f"任务ID: {job_id}")
            check_task_status_example(client, job_id, action_type)
            
    except Exception as e:
        print(f"生成失败: {str(e)}")

def check_task_status_example(client, job_id, action_type=None):
    """示例：查询任务状态"""
    print(f"\n=== 示例：查询任务状态 (ID: {job_id}) ===")
    
    try:
        options = None
        if action_type:
            # 将Submit替换为Query以匹配正确的查询接口类型
            query_action = action_type.replace('Submit', 'Query')
            options = {'action': query_action}
            print(f"使用查询接口类型: {query_action}")
        
        # 查询任务状态
        result = client.get_task_status(job_id, options)
        
        print("任务状态查询结果：")
        print(result)
        
        # 解析任务状态
        if 'Response' in result and 'Status' in result['Response']:
            status = result['Response']['Status']
            print(f"当前任务状态: {status}")
            
            # 根据腾讯云官方文档，任务状态有：WAIT（等待中）、RUN（执行中）、FAIL（任务失败）、DONE（任务成功）
            if status == 'DONE':
                print("任务已成功完成！")
                # 当任务成功时，可以通过ResultFile3Ds获取结果文件
                if 'ResultFile3Ds' in result['Response']:
                    print("3D模型结果文件：")
                    for file_info in result['Response']['ResultFile3Ds']:
                        print(f"- 类型: {file_info.get('Type', '未知')}")
                        print(f"- 下载地址: {file_info.get('Url', '未知')}")
                        if 'PreviewImageUrl' in file_info:
                            print(f"- 预览图地址: {file_info['PreviewImageUrl']}")
            elif status == 'FAIL':
                print("任务失败！")
                # 当任务失败时，可以获取错误信息
                if 'ErrorCode' in result['Response'] and 'ErrorMessage' in result['Response']:
                    print(f"错误代码: {result['Response']['ErrorCode']}")
                    print(f"错误信息: {result['Response']['ErrorMessage']}")
            elif status == 'RUN':
                print("任务正在处理中...")
            elif status == 'WAIT':
                print("任务等待处理...")
            
    except Exception as e:
        print(f"查询失败: {str(e)}")

def monitor_task_until_completion(client, task_id, action_type=None, check_interval=30):
    """示例：监控任务直到完成"""
    print(f"\n=== 示例：监控任务直到完成 (ID: {task_id}) ===")
    
    max_checks = 20  # 最大检查次数
    checks = 0
    
    while checks < max_checks:
        try:
            options = None
            if action_type:
                options = {'action': action_type.replace('Submit', 'Query')}
                
            # 查询任务状态
            result = client.get_task_status(task_id, options)
            
            # 解析任务状态
            if 'Response' in result and 'Status' in result['Response']:
                status = result['Response']['Status']
                print(f"[{time.strftime('%H:%M:%S')}] 任务状态: {status}")
                
                # 如果任务完成或失败，退出循环
                if status in ['DONE', 'FAIL']:
                    print("任务处理完成。")
                    print("最终任务状态：")
                    print(result)
                    return result
                
            # 等待一段时间后再次查询
            time.sleep(check_interval)
            checks += 1
            
        except Exception as e:
            print(f"查询失败: {str(e)}")

def text_to_3d_pro_example():
    """示例：使用专业版接口通过文本生成3D模型"""
    return text_to_3d_example(action_type="SubmitHunyuanTo3DProJob")

def text_to_3d_rapid_example():
    """示例：使用快速版接口通过文本生成3D模型"""
    return text_to_3d_example(action_type="SubmitHunyuanTo3DRapidJob")

def image_to_3d_pro_example():
    """示例：使用专业版接口通过图片生成3D模型"""
    return image_url_to_3d_example(action_type="SubmitHunyuanTo3DProJob")

def image_to_3d_rapid_example():
    """示例：使用快速版接口通过图片生成3D模型"""
    return image_url_to_3d_example(action_type="SubmitHunyuanTo3DRapidJob")

# 如果直接运行此脚本，可以取消下面的注释来测试功能
if __name__ == "__main__":
    print("3D模型生成API使用示例\n==================")
    print("\n支持的接口类型：")
    print("1. SubmitHunyuanTo3DProJob - 专业版接口")
    print("2. SubmitHunyuanTo3DRapidJob - 快速版接口")
    
    # 选择要测试的接口类型
    print("\n请选择要测试的示例：")
    print("1. 默认接口文本生成")
    print("2. 专业版接口文本生成")
    print("3. 快速版接口文本生成")
    print("4. 默认接口图片URL生成")
    print("5. 专业版接口图片URL生成")
    print("6. 快速版接口图片URL生成")
    
    choice = input("请输入选项编号 (1-6): ")
    
    if choice == '1':
        # 测试默认接口文本生成3D模型
        text_result = text_to_3d_example()
    elif choice == '2':
        # 测试专业版接口文本生成3D模型
        text_result = text_to_3d_pro_example()
    elif choice == '3':
        # 测试快速版接口文本生成3D模型
        text_result = text_to_3d_rapid_example()
    elif choice == '4':
        # 测试默认接口图片URL生成3D模型
        image_url_to_3d_example()
    elif choice == '5':
        # 测试专业版接口图片URL生成3D模型
        image_to_3d_pro_example()
    elif choice == '6':
        # 测试快速版接口图片URL生成3D模型
        image_to_3d_rapid_example()
    else:
        print("无效的选项，将运行默认测试...")
        text_result = text_to_3d_example()
    
    # 如果文本生成成功，获取任务ID并查询状态
    # 注意：在实际使用中，任务完成可能需要较长时间
    
    print("\n所有示例运行完毕。")